# include "VideoEncoder.h"
extern "C" {
#include "libavcodec/avcodec.h"
#include "libavformat/avformat.h"
#include "libswscale/swscale.h"
#include "libavdevice/avdevice.h"
#include "libavutil/imgutils.h"
}

# include "EncoderFilter.h"
# include "libyuv.h"

VideoEncoder::VideoEncoder(int _width, int _height, int _fps, int _bit_rate)
{
	avdevice_register_all();
	this->m_video_encode_codec = const_cast<AVCodec*>(avcodec_find_encoder(AV_CODEC_ID_H264));
	if (this->m_video_encode_codec == nullptr) {
		// TODO log

		return;
	}
	// 初始化编码器
	this->m_video_encode_codec_ctx = avcodec_alloc_context3(this->m_video_encode_codec);
	this->m_video_encode_codec_ctx->codec_id = AV_CODEC_ID_H264;
	this->m_video_encode_codec_ctx->codec_type = AVMEDIA_TYPE_VIDEO;
	this->m_video_encode_codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
	this->m_video_encode_codec_ctx->width = _width;
	this->m_video_encode_codec_ctx->height = _height;
	this->m_video_encode_codec_ctx->time_base.num = 1;
	this->m_video_encode_codec_ctx->time_base.den = _fps;
	this->m_video_encode_codec_ctx->max_b_frames = 3;
	this->m_video_encode_codec_ctx->gop_size = 250;
	this->m_video_encode_codec_ctx->bit_rate = _bit_rate * 1000;
	this->m_video_encode_codec_ctx->rc_min_rate = _bit_rate * 1000;
	this->m_video_encode_codec_ctx->rc_max_rate = _bit_rate * 1000;
	this->m_video_encode_codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
	AVDictionary* output_param = nullptr;
	av_dict_set(&output_param, "preset", "veryfast", 0);
	av_dict_set(&output_param, "tune", "zerolatency", 0);
	av_dict_set(&output_param, "profile", "Baseline", 0);
	int r = avcodec_open2(this->m_video_encode_codec_ctx, this->m_video_encode_codec, &output_param);
	av_dict_free(&output_param);
	if (r != 0) {
		// TODO log

		return;
	}
	this->m_video_encode_packet = av_packet_alloc();
	int buf_size = av_image_get_buffer_size(this->m_video_encode_codec_ctx->pix_fmt, this->m_video_encode_codec_ctx->width, this->m_video_encode_codec_ctx->height, 1);
	av_new_packet(this->m_video_encode_packet, buf_size);
	this->m_video_screen_frame = av_frame_alloc();
	this->m_video_encode_frame = av_frame_alloc();

	this->m_video_encode_frame->pts = 0;
	this->m_video_encode_frame->width = _width;
	this->m_video_encode_frame->height = _height;
	this->m_video_encode_frame->format = this->m_video_encode_codec_ctx->pix_fmt;
	av_frame_get_buffer(this->m_video_encode_frame, 0);
	this->m_ready_encode = true;
}

uint8_t* VideoEncoder::getEncodeCodecCtx(size_t* _size) const
{
	if (_size != nullptr) {
		*_size = sizeof(AVCodecContext);
	}
	return reinterpret_cast<uint8_t*>(this->m_video_encode_codec_ctx);
}

uint8_t* VideoEncoder::getExtradata(size_t* _size) const
{
	if (_size != nullptr) {
		*_size = this->m_video_encode_codec_ctx->extradata_size;
	}
	return this->m_video_encode_codec_ctx->extradata;
}

void VideoEncoder::setFilter(EncoderFilter* _filter)
{
	this->m_filter = _filter;
}

void VideoEncoder::onBGRAData(const uint8_t* _data)
{
	if (!this->m_ready_encode) {
		return;
	}
	av_image_fill_arrays(this->m_video_screen_frame->data, this->m_video_screen_frame->linesize, _data, AV_PIX_FMT_BGRA, this->m_video_encode_codec_ctx->width, this->m_video_encode_codec_ctx->height, 1);
	libyuv::ARGBToI420(this->m_video_screen_frame->data[0], this->m_video_screen_frame->linesize[0],
		this->m_video_encode_frame->data[0], this->m_video_encode_frame->linesize[0],
		this->m_video_encode_frame->data[1], this->m_video_encode_frame->linesize[1],
		this->m_video_encode_frame->data[2], this->m_video_encode_frame->linesize[2],
		this->m_video_encode_codec_ctx->width, this->m_video_encode_codec_ctx->height);
	av_frame_unref(this->m_video_screen_frame);
	if (this->m_filter != nullptr) {
		this->m_filter->afterVideoSwr(this->m_video_encode_frame->data[0], this->m_video_encode_frame->nb_samples, this);
	}
	int ret = avcodec_send_frame(this->m_video_encode_codec_ctx, this->m_video_encode_frame);
	ret = avcodec_receive_packet(this->m_video_encode_codec_ctx, this->m_video_encode_packet);
	if (ret != 0) {
		return;
	}
	if (this->m_filter != nullptr) {
		this->m_filter->afterVideoEncode(this->m_video_encode_packet);
	}

	av_packet_unref(this->m_video_encode_packet);
}

VideoEncoder::~VideoEncoder()
{
	this->m_ready_encode = false;

	if (this->m_video_encode_codec_ctx != nullptr) {
		avcodec_send_frame(this->m_video_encode_codec_ctx, nullptr);
		if (this->m_video_encode_packet != nullptr) {
			while (avcodec_receive_packet(this->m_video_encode_codec_ctx, this->m_video_encode_packet) >= 0) {

			}
			av_packet_unref(this->m_video_encode_packet);
		}
	}

	if (this->m_video_encode_packet != nullptr) {
		av_packet_free(&this->m_video_encode_packet);
	}

	if (this->m_video_encode_frame != nullptr) {
		av_frame_free(&this->m_video_encode_frame);
	}

	if (this->m_video_encode_codec_ctx != nullptr) {
		avcodec_free_context(&this->m_video_encode_codec_ctx);
		this->m_video_encode_codec_ctx = nullptr;
	}

	if (this->m_video_screen_frame != nullptr) {
		av_frame_free(&this->m_video_screen_frame);
	}
}
